package org.example;

import com.linkedin.common.urn.DatasetUrn;
import com.linkedin.common.urn.Urn;
import com.linkedin.common.urn.UrnUtils;
import com.linkedin.dataset.DatasetLineageType;
import com.linkedin.metadata.aspect.patch.builder.UpstreamLineagePatchBuilder;
import datahub.client.MetadataWriteResponse;
import datahub.client.rest.RestEmitter;
import org.example.model.ColumnLineage;
import org.example.model.TableLineage;

import java.io.IOException;
import java.net.URISyntaxException;
import java.util.List;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;

import static org.example.DataHubUtil.getDateSet;

public class DataHubClient {

    RestEmitter emitter;
    private String url = "http://192.168.0.13:8080";

    public DataHubClient(String url) {
        this.url = url;
        emitter = RestEmitter.create(b -> b
                .server(url)
        );
    }

    public void sub( List<TableLineage> tableLineages) throws URISyntaxException {

        for (TableLineage tableLineage : tableLineages) {
            UpstreamLineagePatchBuilder upstreamLineagePatchBuilder =         new UpstreamLineagePatchBuilder();
            String env = "PROD";

            Urn tartgetUrn =   UrnUtils.getUrn(getDateSet(tableLineage,env));
            upstreamLineagePatchBuilder.urn(tartgetUrn);
            for (List<ColumnLineage> value : tableLineage.getColumnLineages().values()) {
                for (ColumnLineage columnLineage : value) {
                    try {
                        upstreamLineagePatchBuilder.addUpstream(DatasetUrn.createFromString(
                                getDateSet(columnLineage.getSourceTable(),env)), DatasetLineageType.TRANSFORMED);
                    } catch (URISyntaxException e) {
                        throw new RuntimeException(e);
                    }
                }
            }

            for (String tartFiled : tableLineage.getColumnLineages().keySet()) {
                List<ColumnLineage> value = tableLineage.getColumnLineages().get(tartFiled);
                for (ColumnLineage columnLineage : value) {
                    upstreamLineagePatchBuilder.addFineGrainedUpstreamField(new Urn("urn:li:schemaField:("+getDateSet(columnLineage.getSourceTable(),env)+","+columnLineage.getSourceColumn()+")"),
                            0.5f,"dddd",new Urn("urn:li:schemaField:("+getDateSet(tableLineage,env)+","+tartFiled+")"),null);
                }


            }

            Future<MetadataWriteResponse> response = null;
            try {

                response = emitter.emit(upstreamLineagePatchBuilder.build());
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
            try {
                System.out.println(response.get().getResponseContent());
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            } catch (ExecutionException e) {
                throw new RuntimeException(e);
            }


        }
    }
}
